"""
创建问答json文件
{
    "问题1":{
        "主体":["主体1","主体3","主体3"..],
        "问题1分词后的句子":["word1","word2","word3"...],
        "答案":"答案"
    },
    "问题2":{
        ...
    }
}
"""
import json
import pandas as pd

from lib import cut
import config

# 问题文件路径
q_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目\案例-chat_service\corpus" \
         r"\dnn\recall/Q.txt"
# 答案文件路劲
a_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目\案例-chat_service\corpus" \
         r"\dnn\recall/A.txt"
# excel文件路径
excel_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目\案例-chat_service\corpus" \
             r"\dnn\recall/excel.xlsx"


def process_recall_corpus():
    q_data = open(q_path, mode="r", encoding="UTF-8").readlines()
    a_data = open(a_path, mode="r", encoding="UTF-8").readlines()

    qa_dict = {}
    for q, a in zip(q_data, a_data):
        q = q.strip()
        qa_dict[q] = {}
        qa_dict[q]["answer"] = a.strip()
        qa_dict[q]["cut_by_word"] = cut(q, by_word=True)
        result = cut(q, with_sg=True)
        # print(result)
        qa_dict[q]["cut"] = [i[0].strip() for i in result]
        entity = [i[0].strip() for i in result if i[1].strip() == "kc"]
        qa_dict[q]["entity"] = entity
        # break

    df = pd.read_excel(excel_path)
    # print(df)
    for q, a in zip(df['问题'], df['答案']):
        # print(q, a)
        q = q.strip()
        qa_dict[q] = {}
        qa_dict[q]["answer"] = a.strip()
        qa_dict[q]["cut_by_word"] = cut(q, by_word=True)
        lines_cut_with_sg = cut(q, by_word=False, with_sg=True)
        qa_dict[q]["cut"] = [i[0].strip() for i in lines_cut_with_sg]
        entity = [i[0].strip() for i in lines_cut_with_sg if i[1] == "kc"]
        qa_dict[q]["entity"] = entity

    # # 获取保存数据路径
    # save_data_path = config.recall_corpus_bm25_path if method.lower() == "bm25" \
    #     else config.recall_corpus_tfidf_path

    json.dump(qa_dict, open(config.recall_corpus_tfidf_path,
                            "w", encoding="UTF-8"),
              ensure_ascii=False, indent=2)
